{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE FlexibleContexts      #-}

module Language.Haskell.Liquid.Constraint.Qualifier
  ( giQuals
  , useSpcQuals
  )
  where

import           Prelude hiding (error)
import           Data.List                (delete, nub, partition)
import           Data.Maybe               (isJust, catMaybes, fromMaybe, isNothing)
import qualified Data.HashSet        as S
import qualified Data.HashMap.Strict as M
import           Debug.Trace (trace)
import           Control.Monad.Reader

import           Language.Fixpoint.Types                  hiding (panic, mkQual)
import qualified Language.Fixpoint.Types.Config as FC
import           Language.Fixpoint.Misc                  (fst3)
import           Language.Fixpoint.SortCheck
import           Language.Haskell.Liquid.Types.RefType
import           Language.Haskell.Liquid.GHC.Misc         (getSourcePos)
import           Language.Haskell.Liquid.Misc             (condNull)
import           Language.Haskell.Liquid.Types.PredType
import           Liquid.GHC.API hiding (Expr, mkQual, panic)

import           Language.Haskell.Liquid.Types.Errors
import           Language.Haskell.Liquid.Types.RType
import           Language.Haskell.Liquid.Types.RTypeOp
import           Language.Haskell.Liquid.Types.Specs
import           Language.Haskell.Liquid.Types.Types
import           Language.Haskell.Liquid.UX.Config
import Language.Fixpoint.Types.Config (ElabFlags)


--------------------------------------------------------------------------------
giQuals :: TargetInfo -> SEnv Sort -> [Qualifier]
--------------------------------------------------------------------------------
giQuals info lEnv
  =  notracepp ("GI-QUALS: " ++ showpp lEnv)
  $  currentSpec
  ++ condNull (useSpcQuals info) otherSpecs
  ++ condNull (useSigQuals info) (sigQualifiers  info lEnv)
  ++ condNull (useAlsQuals info) (alsQualifiers  info lEnv)
  where 
    (currentSpec, otherSpecs) = partition isQualifierFromCurrentModule (gsQualifiers . gsQual . giSpec $ info)
    isQualifierFromCurrentModule qual = 
      fst3 (sourcePosElts (qPos qual)) == giTarget (giSrc info)


-- --------------------------------------------------------------------------------
-- qualifiers :: GhcInfo -> SEnv Sort -> [Qualifier]
-- --------------------------------------------------------------------------------
-- qualifiers info env = spcQs ++ genQs
  -- where
    -- spcQs           = gsQualifiers spc
    -- genQs           = specificationQualifiers info env
    -- n               = maxParams (getConfig spc)
    -- spc             = spec info

maxQualParams :: (HasConfig t) => t -> Int
maxQualParams = maxParams . getConfig

-- | Use explicitly given qualifiers
useSpcQuals :: (HasConfig t) => t -> Bool
useSpcQuals i = useQuals i && not (useAlsQuals i)

-- | Scrape qualifiers from function signatures (incr :: x:Int -> {v:Int | v > x})
useSigQuals :: (HasConfig t) => t -> Bool
useSigQuals i = useQuals i && not (useAlsQuals i)

-- | Scrape qualifiers from refinement type aliases (type Nat = {v:Int | 0 <= 0})
useAlsQuals :: (HasConfig t) => t -> Bool
useAlsQuals i = useQuals i && i `hasOpt` higherOrderFlag && not (needQuals i)

useQuals :: (HasConfig t) => t -> Bool
useQuals = (FC.All /=) . eliminate . getConfig

needQuals :: (HasConfig t) => t -> Bool
needQuals = (FC.None == ) . eliminate . getConfig

--------------------------------------------------------------------------------
alsQualifiers :: TargetInfo -> SEnv Sort -> [Qualifier]
--------------------------------------------------------------------------------
alsQualifiers info lEnv
  = [ q | a <- gsRTAliases . gsQual . giSpec $ info
        , q <- refTypeQuals lEnv ef (loc a) tce (rtBody (val a))
        , length (qParams q) <= k + 1
        , validQual lEnv q
    ]
    where
      ef  = maybe (FC.ElabFlags False) FC.solverFlags . smtsolver . getConfig $ info
      k   = maxQualParams info
      tce = gsTcEmbeds . gsName . giSpec $ info

validQual :: SEnv Sort -> Qualifier -> Bool
validQual lEnv q = isJust $ checkSortExpr (srcSpan q) env (qBody q)
  where
    env          = unionSEnv lEnv qEnv
    qEnv         = M.fromList (qualBinds q)


--------------------------------------------------------------------------------
sigQualifiers :: TargetInfo -> SEnv Sort -> [Qualifier]
--------------------------------------------------------------------------------
sigQualifiers info lEnv
  = [ q | (x, t) <- specBinders info
        , not (x `S.member` S.fromList (specAxiomVars info))
        , q <- refTypeQuals lEnv ef (getSourcePos x) tce (val t)
        -- NOTE: large qualifiers are VERY expensive, so we only mine
        -- qualifiers up to a given size, controlled with --max-params
        , length (qParams q) <= k + 1
    ]
    where
      ef  = maybe (FC.ElabFlags False) FC.solverFlags . smtsolver . getConfig $ info
      k   = maxQualParams info
      tce = gsTcEmbeds . gsName . giSpec $ info

specBinders :: TargetInfo -> [(Var, LocSpecType)]
specBinders info = mconcat
  [ gsTySigs  (gsSig  sp)
  , gsAsmSigs (gsSig  sp)
  , gsRefSigs (gsSig  sp)
  , gsCtors   (gsData sp)
  , if info `hasOpt` scrapeInternals then gsInSigs (gsSig sp) else []
  ]
  where
    sp  = giSpec info

specAxiomVars :: TargetInfo -> [Var]
specAxiomVars =  gsReflects . gsRefl . giSpec

-- GRAVEYARD: scraping quals from imports kills the system with too much crap
-- specificationQualifiers info = {- filter okQual -} qs
--   where
--     qs                       = concatMap refTypeQualifiers ts
--     refTypeQualifiers        = refTypeQuals $ tcEmbeds spc
--     ts                       = val <$> t1s ++ t2s
--     t1s                      = [t | (x, t) <- tySigs spc, x `S.member` definedVars]
--     t2s                      = [] -- [t | (_, t) <- ctor spc                            ]
--     definedVars              = S.fromList $ defVars info
--     spc                      = spec info
--
-- okQual                       = not . any isPred . map snd . q_params
--   where
--     isPred (FApp tc _)       = tc == stringFTycon "Pred"
--     isPred _                 = False


-- TODO: rewrite using foldReft'
--------------------------------------------------------------------------------
refTypeQuals :: SEnv Sort -> ElabFlags -> SourcePos -> TCEmb TyCon -> SpecType -> [Qualifier]
--------------------------------------------------------------------------------
refTypeQuals lEnv ef l tce t0    = go emptySEnv t0
  where
    scrape                    = refTopQuals lEnv ef l tce t0
    add x t γ                 = insertSEnv x (rTypeSort tce t) γ
    goBind x t γ t'           = go (add x t γ) t'
    go γ t@(RVar _ _)         = scrape γ t
    go γ (RAllT _ t _)        = go γ t
    go γ (RAllP p t)          = go (insertSEnv (pname p) (rTypeSort tce (pvarRType p :: RSort)) γ) t
    go γ t@(RAppTy t1 t2 _)   = go γ t1 ++ go γ t2 ++ scrape γ t
    go γ (RFun x _ t t' _)    = go γ t ++ goBind x t γ t'
    go γ t@(RApp c ts rs _)   = scrape γ t ++ concatMap (go γ') ts ++ goRefs c γ' rs
                                where γ' = add (rTypeValueVar t) t γ
    go γ (RAllE x t t')       = go γ t ++ goBind x t γ t'
    go γ (REx x t t')         = go γ t ++ goBind x t γ t'
    go _ _                    = []
    goRefs c g rs             = concat $ zipWith (goRef g) rs (rTyConPVs c)
    goRef _ (RProp _ (RHole _)) _ = []
    goRef g (RProp s t)  _    = go (insertsSEnv' g s) t
    insertsSEnv'              = foldr (\(x, t) γ -> insertSEnv x (rTypeSort tce t) γ)


refTopQuals :: (PPrint t, Reftable t, SubsTy RTyVar RSort t, Reftable (RTProp RTyCon RTyVar (UReft t)))
            => SEnv Sort
            -> ElabFlags
            -> SourcePos
            -> TCEmb TyCon
            -> RType RTyCon RTyVar r
            -> SEnv Sort
            -> RRType (UReft t)
            -> [Qualifier]
refTopQuals lEnv ef l tce t0 γ rrt
  = [ mkQ' v so pa  | let (RR so (Reft (v, ra))) = rTypeSortedReft tce rrt
                   , pa                        <- conjuncts ra
                   , not $ isHole    pa
                   , notracepp ("refTopQuals: " ++ showpp pa)
                     $ isNothing $ runReader (checkSorted (srcSpan l) (insertSEnv v so γ') pa) ef
    ]
    ++
    [ mkP s e | let (MkUReft _ (Pr ps)) = fromMaybe (msg rrt) $ stripRTypeBase rrt
              , p                      <- findPVar (ty_preds $ toRTypeRep t0) <$> ps
              , (s, _, e)              <- pargs p
    ]
    where
      mkQ'  = mkQual  lEnv l     t0 γ
      mkP   = mkPQual lEnv l tce t0 γ
      msg t = panic Nothing $ "Qualifier.refTopQuals: no typebase" ++ showpp t
      γ'    = unionSEnv' γ lEnv

mkPQual :: (PPrint r, Reftable r, SubsTy RTyVar RSort r, Reftable (RTProp RTyCon RTyVar r))
        => SEnv Sort
        -> SourcePos
        -> TCEmb TyCon
        -> t
        -> SEnv Sort
        -> RRType r
        -> Expr
        -> Qualifier
mkPQual lEnv l tce t0 γ t e = mkQual lEnv l t0 γ' v so pa
  where
    v                      = "vv"
    so                     = rTypeSort tce t
    γ'                     = insertSEnv v so γ
    pa                     = PAtom Eq (EVar v) e

mkQual :: SEnv Sort
       -> SourcePos
       -> t
       -> SEnv Sort
       -> Symbol
       -> Sort
       -> Expr
       -> Qualifier
mkQual lEnv l _ γ v so p   = mkQ "Auto" ((v, so) : xts) p l
  where
    xs   = delete v $ nub $ syms p
    xts  = catMaybes $ zipWith (envSort l lEnv γ) xs [0..]

envSort :: SourcePos -> SEnv Sort -> SEnv Sort -> Symbol -> Integer -> Maybe (Symbol, Sort)
envSort l lEnv tEnv x i
  | Just t <- lookupSEnv x tEnv = Just (x, t)
  | Just _ <- lookupSEnv x lEnv = Nothing
  | otherwise                   = Just (x, ai)
  where
    ai             = trace msg $ fObj $ Loc l l $ tempSymbol "LHTV" i
    msg            = "Unknown symbol in qualifier: " ++ show x
